#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "../hash_map/hash_map.h"

#define max(a, b) (((a) > (b)) ? (a) : (b))

#define L_SLAB_SIZE 65536
#define G_SLAB_SIZE 262144 // devouring memory

static const int CACHE_NUMBER = 10;

#define BLOCKS_MAX 3000
#define SLABS_MAX  200

static size_t current_sa_size;
static hash_map_t hm_slab;

static struct {
	char *mem; // ptr to the beginnig of available mem
	void *ptr_to_del;
	char *mem_bound;
} sa_allocator;

typedef struct block {
	char *ptr;
	size_t size;
	struct block *next;
} block_t;

typedef struct slab {
	char *ptr;
	block_t *free_blocks;
	int block_num;
	struct slab *prev;
	struct slab *next;
	int used;
} slab_t;

typedef struct cache {
	int slab_size;
	int block_size;
	slab_t *partial_slabs;
	slab_t *empty_slabs;
	slab_t *full_slabs;
} cache_t;

static int cur_slab = 0;
static int cur_block = 0;

static slab_t slabs[SLABS_MAX];
static block_t blocks[BLOCKS_MAX];

static cache_t caches[] = {
	{ L_SLAB_SIZE, 128,   NULL, NULL, NULL }, // ~512 variables
	{ L_SLAB_SIZE, 256,   NULL, NULL, NULL }, // ~256 variables
	{ L_SLAB_SIZE, 512,   NULL, NULL, NULL }, // ~128 variables
	{ L_SLAB_SIZE, 1024,  NULL, NULL, NULL }, // ~64 variables
	{ L_SLAB_SIZE, 2048,  NULL, NULL, NULL }, // ~32 variables
	{ G_SLAB_SIZE, 4096,  NULL, NULL, NULL }, // ~64 variables
	{ G_SLAB_SIZE, 8192,  NULL, NULL, NULL }, // ~32 variables
	{ G_SLAB_SIZE, 16384, NULL, NULL, NULL }, // ~16 variables
	{ G_SLAB_SIZE, 32768, NULL, NULL, NULL }, // ~8 variables
	{ G_SLAB_SIZE, 65536, NULL, NULL, NULL }, // ~4 variables
};

static void poison_mem(char *mem, size_t sz) {
	int x = 0xDEADBEAF, i;
	for (i = 0; i < sz / sizeof(int); ++i) {
		memcpy(mem, &x, sizeof(int));
		mem += 4;
	}
	memcpy(mem, &x, sz % sizeof(int));
}

int sa_init(size_t size) {
	int i, j;
	block_t *block;
	slab_t *slab;

	current_sa_size = 1638400;
	size_t sz = max(current_sa_size, size);
	char *mem = (char *)malloc(sz);
	if (!mem)
		return 0;
#ifdef _POISON
	poison_mem(mem, sz);
#endif
	sa_allocator.ptr_to_del = mem;
	sa_allocator.mem_bound = mem + sz;
	map_init(&hm_slab);
	memset(slabs, 0, SLABS_MAX * sizeof(slab_t));
	memset(blocks, 0, BLOCKS_MAX * sizeof(block_t));

	for (i = 0; i < CACHE_NUMBER; ++i) {
		slab = &slabs[cur_slab++];
		slab->next = caches[i].partial_slabs;
		caches[i].partial_slabs = slab;
		slab->ptr = mem;
		slab->block_num = caches[i].slab_size / caches[i].block_size;

		for (j = 0; j < slab->block_num; ++j) {
			block = &blocks[cur_block++];
			block->ptr = mem;
			block->size = caches[i].block_size;
			mem += caches[i].block_size;
			block->next = slab->free_blocks;
			slab->free_blocks = block;
		}
	}
	sa_allocator.mem = mem;
	return 1;
}

void *sa_alloc(size_t sz) {
	if (sz <= 128)
		sz = 128;
	else if (sz <= 512)
		sz = 512;
	else if (sz <= 1024)
		sz = 1024;
	else if (sz <= 2048)
		sz = 2048;
	else if (sz <= 4096)
		sz = 4096;
	else if (sz <= 8192)
		sz =  8192;
	else if (sz <= 16384)
		sz = 16384;
	else if (sz <= 32768)
		sz = 32768;
	else
		sz = 65536;

	int i, idx = log2(sz >> 7);
	slab_t *slab = caches[idx].partial_slabs;

	if (!slab)
		return NULL;
	block_t *block = slab->free_blocks,
					*tblock;
	if (!block)
		return NULL;
	if (!map_insert(&hm_slab, (size_t)block->ptr, block))
		return NULL;

	slab->free_blocks = block->next;
	++slab->used;
	block->size = sz;
	if (slab->free_blocks) {
		return block->ptr;
	}

	if (!slab->free_blocks) {
		caches[idx].partial_slabs = caches[idx].partial_slabs->next;
		// add to full slabs
		if (caches[idx].full_slabs)
			caches[idx].full_slabs->prev = slab;

		slab->next = caches[idx].full_slabs;
		caches[idx].full_slabs = slab;
		// delete from partial slabs
		if (caches[idx].partial_slabs)
			return block->ptr;
		// get from empty slabs
		if (caches[idx].empty_slabs) {
			caches[idx].partial_slabs = caches[idx].empty_slabs;
			caches[idx].empty_slabs = caches[idx].empty_slabs->next;
			return block->ptr;
		}
		int findx = (caches[idx].slab_size == L_SLAB_SIZE) ? 0 : 5;
		for (i = 0; i < findx + 5; ++i) {
			if (caches[i].empty_slabs) {
				caches[idx].partial_slabs = caches[i].empty_slabs;
				caches[i].empty_slabs = caches[i].empty_slabs->next;
				return block->ptr;
			}
		}
	}
	// if we are here - we have to create new slab
	if ((sa_allocator.mem_bound < sa_allocator.mem + caches[idx].slab_size) ||
		(cur_slab >= SLABS_MAX) || (cur_block + caches[idx].slab_size/caches[idx].block_size >= BLOCKS_MAX))
		return block->ptr;
	slab = &slabs[cur_slab++];
	slab->block_num = caches[idx].slab_size / caches[idx].block_size;
	current_sa_size += caches[idx].slab_size;
	slab->next = caches[idx].partial_slabs;
	caches[idx].partial_slabs = slab;
	caches[idx].partial_slabs->ptr = sa_allocator.mem;

	// Divide slab into blocks
	for (i = 0; i < slab->block_num; ++i) {
		tblock = &blocks[cur_block++];
		tblock->ptr = sa_allocator.mem;
		tblock->size = sz;
		sa_allocator.mem += caches[idx].block_size;
		tblock->next = slab->free_blocks;
		slab->free_blocks = tblock;
	}
	return block->ptr;
}

void sa_free(void *ptr_) {
	block_t *block = (block_t *)map_delete(&hm_slab, (size_t)ptr_);
	int idx = log2(block->size >> 7),
			slab_found = 0, sz = block->size;
	cache_t *cache = &caches[idx];
	slab_t *slab;
	char *ptr = (char *)ptr_;
	int i;

	for (slab = cache->full_slabs; slab; slab = slab->next) {
		if ((slab->ptr <= ptr) && (slab->ptr + caches[idx].slab_size >= ptr + block->size)) {
			if (slab->prev)
				slab->prev->next = slab->next;
			else
				cache->full_slabs = cache->full_slabs->next;
			slab->next = caches[idx].partial_slabs;
			caches[idx].partial_slabs = slab;
			slab_found = 1;
			break;
		}
	}
	if (!slab_found) {
		for (slab = cache->partial_slabs; slab; slab = slab->next) {
			if ((slab->ptr <= ptr) && (slab->ptr + caches[idx].slab_size >= ptr + block->size)) {
				slab_found = 1;
				break;
			}
		}
	}
	block->next = slab->free_blocks;
	slab->free_blocks = block;
	--slab->used;

	if ((!slab->used) && (!slab->next)) {
		// partial -> empty
		slab->next = caches[idx].empty_slabs;
		caches[idx].empty_slabs = slab;
	}
#ifdef _POISON
	poison_mem(ptr_, sz);
#endif
}

void sa_deinit() {
	int i;
	free(sa_allocator.ptr_to_del);
	for (i = 0; i < CACHE_NUMBER; ++i) {
		caches[i].empty_slabs = NULL;
		caches[i].partial_slabs = NULL;
		caches[i].full_slabs = NULL;
	}
	cur_block = 0;
	cur_slab = 0;
}
